import torch  # 确保导入 torch
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info
from zmq_server import ImageZmqServer
import threading
import time
import os
import re
import json

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    "Qwen/Qwen2.5-VL-7B-Instruct", torch_dtype="auto", device_map=None
).to(device)

processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")

def process_image(image_path, prompt_text=""):
    """处理图像并返回描述结果"""
    messages = [        
        {
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "image": image_path,
                },
                {"type": "text", "text": prompt_text},
            ],
        }
    ]

    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    )
    inputs = inputs.to(device)

    generated_ids = model.generate(**inputs, max_new_tokens=512)
    generated_ids_trimmed = [
        out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    output_text = processor.batch_decode(
        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )
    return output_text

def extract_all_bbox(text):
    """从文本中提取所有bbox_2d坐标"""
    # 检查text是否为列表，如果是则连接为一个字符串
    if isinstance(text, list):
        text = '\n'.join(text)
    
    bbox_list = []
    
    # 使用正则表达式查找所有JSON格式的bbox_2d
    pattern = r'"bbox_2d"\s*:\s*\[\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*\]\s*,\s*"label"\s*:\s*"([^"]+)"'
    matches = re.findall(pattern, text)
    
    for match in matches:
        # 提取四个坐标值和标签
        x1, y1, x2, y2, label = match
        bbox_list.append([int(x1), int(y1), int(x2), int(y2), label])
    
    return bbox_list if bbox_list else None

def start_zmq_server():
    """启动ZMQ服务器线程"""
    save_dir = "output"
    os.makedirs(save_dir, exist_ok=True)
    
    server = ImageZmqServer('0.0.0.0', 5555, save_dir=save_dir)
    
    # 创建一个监控线程，检查是否有新图像和文本
    def monitor_images():
        last_image_time = None
        prompt_text = ""  
        
        while True:
            # 检查是否有新图像
            if server.latest_timestamp != last_image_time and server.latest_image is not None:
                print("检测到新图像，开始处理...")
                last_image_time = server.latest_timestamp
                # 检查是否有新文本
                if server.latest_text is not None and server.latest_text_timestamp is not None:
                    prompt_text = server.latest_text
                    print(f"接收到新的提示文本: {prompt_text}")
                
                    # 使用保存的图像路径
                    image_path = f"{save_dir}/receive_image.jpg"
                    if os.path.exists(image_path):
                        # 处理图像并获取描述，使用最新的提示文本
                        description_list = process_image(image_path, prompt_text)
                        print(f"图像描述结果: {description_list}")
                            
                        # 检查描述中是否包含bbox坐标
                        bbox_list = extract_all_bbox(description_list)
                            
                        if bbox_list:
                            # 如果找到了bbox，通过ZMQ发送回去
                            print(f"检测到{len(bbox_list)}个bbox:")
                            for i, bbox in enumerate(bbox_list):
                                print(f"  bbox {i+1}: {bbox}")
                            
                            # 使用服务器对象的send_coords方法发送所有坐标，标志位为True
                            server.send_coords(bbox_list, found_bbox=True)
                        else:
                            print("未检测到bbox坐标")
                            # 发送坐标为0的默认值，标志位为False
                            server.send_coords([[0, 0, 0, 0, "无"]], found_bbox=False)
                
            time.sleep(1)  # 每秒检查一次
    
    # 启动监控线程
    monitor_thread = threading.Thread(target=monitor_images, daemon=True)
    monitor_thread.start()
    
    # 运行服务器（这会阻塞当前线程）
    server.run()

if __name__ == "__main__":
    # 启动ZMQ服务器
    start_zmq_server()